#!/usr/bin/env python3

# Copyright © 2025 Wenze Wei
#
# This file is part of Pisces L1.
#
# Licensed under the Creative Commons Attribution-NonCommercial 4.0 International License (CC BY-NC 4.0).
# You may not use this file except in compliance with the License.
# Commercial use is strictly prohibited.
# You may obtain a copy of the License at
#
#     https://creativecommons.org/licenses/by-nc/4.0/
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
import argparse
from pathlib import Path

# Get the absolute path of the directory where the current script is located
ROOT = os.path.abspath(os.path.dirname(__file__))

# List of available commands with their descriptions
COMMANDS = [
    'setup',      # Environment setup
    'source',     # Activate virtual environment
    'pull',       # Pull latest code from remote
    'train',      # Train model
    'infer',      # Inference
    'check',      # Check GPU/deps
    'monitor',    # System monitor
    'download',   # Download datasets
    'arrow',      # Arrow/JSON conversion
    'help',       # Show help for commands
]

def main():
    """
    Main function to handle command-line arguments and execute corresponding commands.
    This function parses the command-line arguments and calls the appropriate function
    based on the specified command.
    """
    # Create an argument parser for the Pisces L1 Management Tool
    parser = argparse.ArgumentParser(description="Pisces L1 Management Tool (manage.py)")
    # Add a positional argument for the command to execute
    parser.add_argument('command', nargs='?', choices=COMMANDS, help="Command to execute")
    # Add an argument for the checkpoint file used in inference
    parser.add_argument('--ckpt', default='', help='Checkpoint file (for infer)')
    # Add an argument for the prompt used in inference
    parser.add_argument('--prompt', default='Hello, please introduce yourself', help='Prompt (for infer)')
    # Add an argument for the image path used in inference
    parser.add_argument('--image', default='', help='Image path (for infer)')
    # Add an argument for the maximum number of samples per dataset when downloading
    parser.add_argument('--max_samples', type=int, default=50000, help='Max samples per dataset (for download)')
    # Add an argument for the directory containing .json files to merge into one .arrow file
    parser.add_argument('--json_dir', default='', help='[arrow] Directory containing .json files to merge into one .arrow')
    # Add an argument for the output .arrow file path
    parser.add_argument('--arrow_out', default='', help='[arrow] Output .arrow file path')
    # Add an argument for the input .arrow file path to convert to .json
    parser.add_argument('--arrow_in', default='', help='[arrow] Input .arrow file path to convert to .json')
    # Add an argument for the output .json file path (single file)
    parser.add_argument('--json_out', default='', help='[arrow] Output .json file path (single file)')
    # Parse the known arguments and leave the extra arguments unparsed
    args, extra = parser.parse_known_args()
    
    if args.command is None or args.command == 'help':
        # Import and call the help function if no command is specified or the command is 'help'
        from tools.help import help
        help()
    elif args.command == 'train':
        # Import functions to add training arguments and start training
        from tools.train import add_train_args, train
        parser = add_train_args(parser)
        args = parser.parse_args()
        train(args)
    elif args.command == 'infer':
        # Import and call the inference function
        from tools.infer import infer
        infer(args)
    elif args.command == 'check':
        # Import and call the check function
        from tools.check import check
        check(args)
    elif args.command == 'monitor':
        # Import and call the monitor function
        from tools.monitor import monitor
        monitor()
    elif args.command == 'download':
        # Import functions to download and optimize datasets
        from data.download import download_datasets, optimize_datasets
        download_datasets(args.max_samples)
        optimize_datasets(max_keep=5000)
    elif args.command == 'arrow':
        # Import and call the arrow conversion function
        from tools.arrow import arrow
        arrow(args)
    elif args.command == 'setup':
        # Import and call the setup function
        from tools.setup import setup
        setup(args)
    elif args.command == 'source':
        # Import necessary modules for environment activation
        import os
        import sys
        # Check if the current operating system is Windows
        is_windows = os.name == 'nt'
        # Get the path to the virtual environment directory
        venv_dir = os.path.join(os.path.dirname(__file__), 'pieces_env')
        
        if is_windows:
            # Get the Windows shell executable
            shell = os.environ.get("COMSPEC", "cmd.exe")
            # Check if using PowerShell
            if "powershell.exe" in shell.lower() or "pwsh.exe" in shell.lower():
                # Get the path to the PowerShell activation script
                activate = os.path.join(venv_dir, "Scripts", "Activate.ps1")
                print("✅\tAuto-entering Pisces venv shell (PowerShell)...")
                os.execv(shell, [shell, "-NoExit", "-Command", f". '{activate}'"])
            else:
                # Get the path to the cmd activation script
                activate = os.path.join(venv_dir, "Scripts", "activate.bat")
                print("✅\tAuto-entering Pisces venv shell (Windows cmd)...")
                os.execv(shell, [shell, "/K", activate])
        else:
            # Get the Unix-like shell executable
            shell = os.environ.get("SHELL", "/bin/bash")
            # Get the path to the activation script
            activate = os.path.join(venv_dir, "bin", "activate")
            print("✅\tAuto-entering Pisces venv shell (Linux/Mac)...")
            os.execv(shell, [shell, "-i", "-c", f"source '{activate}'; exec {shell}"])
    elif args.command == 'pull':
        # Import necessary modules for git operations
        import subprocess
        import sys
        # Define the remote repository URL
        remote_url = 'https://gitee.com/dunimd/piscesl1.git'
        try:
            print(f"🟧\tPulling latest code from {remote_url}...")
            # Fetch all remote branches
            subprocess.run(['git', 'fetch', '--all'], check=True)
            # Reset the local branch to the latest remote master
            subprocess.run(['git', 'reset', '--hard', 'origin/master'], check=True)
            print("✅\tCode successfully updated to the latest version")
        except subprocess.CalledProcessError as e:
            print(f"❌\tFailed to pull code: {e}")
            sys.exit(1)
    elif args.command == 'help':
        # Import and call the help function if the command is 'help'
        from tools.help import help
        help()
    else:
        # Print an error message and exit if the command is unknown
        print(f"❌\tUnknown command: {args.command}")
        sys.exit(1)

if __name__ == "__main__":
    # Call the main function if the script is run directly
    main()